#ifndef BOOTSTRAPPING_SLOTS_H
#define BOOTSTRAPPING_SLOTS_H
#include <vector>
#include <cassert>
#include <ostream>
#include <memory>
#include "polyarith.h"

// 包含有使用明文空间槽的操作
extern size_t kswitchCounter;

void LogRelin();
void LogGalois();

inline uint64_t Log2Ceil(uint64_t x)
{
    if (x == 0) {
        return 0;
    }
    unsigned long result = 0;
    while (x != 0) {
        x = x >> 1;
        result += 1;
    }
    return result - 1;
};

inline uint64_t Log2Exact(uint64_t x)
{
    unsigned long result = Log2Ceil(x);
    if (((uint64_t)1 << result) == x) {
        return result;
    } else {
        throw std::invalid_argument("Not a power of 2");
    }
};

inline uint64_t HighestDividingPower2(uint64_t x)
{
    if (x == 0) {
        return std::numeric_limits<uint64_t>::digits;
    }
    unsigned long result = 0;
    while ((x & 1) == 0) {
        result += 1;
        x = x >> 1;
    }
    return result;
};

void TestGAutomorphisms();

// 表示环的分解结构，n为槽的个数
class SlotRing {
    // 槽的个数
    uint64_t n;
    // 槽的秩
    uint64_t d;
    // 在Z/p^eZ中的素数p
    uint64_t p;
    // 在Z/p^eZ中的素数指数e
    uint64_t e;

    // 自同构群的第一个生成器
    uint64_t g1;

    // 自同构群的第二个生成器
    uint64_t g2;

    // 槽组是否循环
    bool slotGroupCyclic;

    // 槽多项式是否稀疏
    bool dSparse;

    // 以2为底2N的对数
    uint64_t log2N;

    std::tuple<size_t, size_t> pLog;

    // 环的模数
    uint64_t modPe;

    // 下标环的模数
    uint64_t mod2N;

    // X^N + 1多项式环的模数
    PolyModulus mod2NCyclotomic;

    // 一个槽的生成元的最小多项式；目前，这与第一个槽匹配，即 f 模 slotModulus 是 f 的槽 0 部分
    PolyModulus slotModulus;

    // 存储的槽值为1或0所对应的编码多项式环中的元素
    std::vector<poly> unitVectors;

    void InitSlotGroup();
    void InitPLog();
    void InitSlotModulus(poly baseUnitVector);

    void ApplyFrobenius(poly &x, size_t iters, const uint64_t *mod = nullptr) const;

    // 计算以g1为底的自同构，如果mod为nullptr，mod为p^e
    void G1Automorphism(poly &x, size_t iters, const uint64_t *mod = nullptr) const;

    // 计算以g2为底的自同构，如果mod为nullptr，mod为p^e
    void G2Automorphism(poly &x, const uint64_t *mod = nullptr) const;

public:
    /*
     * 表示主环的子环，只要用于单槽
     */
    struct SubringView {
        const uint64_t &scalarMod;
        const PolyModulus &polyMod;

        poly Generator();
    };

    // 用map封装槽的旋转
    class Rotation {
        // fMask为1的元素“向前”移动s
        poly fMask;
        // bMask为1的元素“向后”移动s
        poly bMask;
        size_t s;
        const SlotRing &slotRing;
        size_t blockSize;
        mutable std::unique_ptr<Plaintext> fMaskPlain;
        mutable std::unique_ptr<Plaintext> bMaskPlain;

        inline Rotation(poly fMask, poly bMask, size_t s, const SlotRing &slotRing, size_t blockSize)
            : fMask(std::move(fMask)),
              bMask(std::move(bMask)),
              s(s),
              slotRing(slotRing),
              blockSize(blockSize),
              fMaskPlain(nullptr),
              bMaskPlain(nullptr)
        {}

        size_t EffectiveBlockSize() const;

    public:
        Rotation(const Rotation &) = default;
        Rotation(Rotation &&) = default;
        ~Rotation() = default;

        poly operator () (const poly &x) const;
        void ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk, Ciphertext &result);

        std::tuple<uint32_t, uint32_t> GaloisElements() const;

        inline const poly &GetForwardMask() const
        {
            return fMask;
        }
        inline const poly &GetBackwardMask() const
        {
            return bMask;
        }
        const std::tuple<uint64_t, uint64_t> GetForwardG1G2Decomp() const;
        const std::tuple<uint64_t, uint64_t> GetBackwardG1G2Decomp() const;

        friend SlotRing;
    };

    /*
     * 封装逐槽Frobenius_ 映射
     */
    class Frobenius {
        size_t iters;
        const SlotRing &slotRing;

        inline Frobenius(size_t iters, const SlotRing &slotRing) : iters(iters), slotRing(slotRing) {}

    public:
        Frobenius(const Frobenius &) = default;
        Frobenius(Frobenius &&) = default;
        ~Frobenius() = default;

        poly operator () (poly x) const;

        uint32_t galois_element() const;

        const std::tuple<uint64_t, uint64_t> GetG1G2Decomp() const;

        friend SlotRing;
    };

    /*
     * 封装了由相应的环外有理数域的伽罗瓦自同构（Galois Automorphism）所引起的基本环的原始自同构（Automorphism）。
     */
    class RawAuto {
        // 存储自同构需要的g1，g2
        std::tuple<size_t, size_t> g1g2Decomp;
        const SlotRing &slotRing;

        inline RawAuto(const SlotRing &slotRing, std::tuple<size_t, size_t> g1G2Decomp)
            : slotRing(slotRing), g1g2Decomp(g1G2Decomp)
        {}

    public:
        RawAuto(const RawAuto &) = default;
        RawAuto(RawAuto &&) = default;
        ~RawAuto() = default;

        poly operator () (poly x) const;
        void ApplyCiphertext(Ciphertext &in, Evaluator &eval, GaloisKeys &gk, Ciphertext &result);
        uint32_t galois_element() const;
        bool is_identity() const;

        friend SlotRing;
    };

    SlotRing(uint64_t log2N, uint64_t p, uint64_t e, poly baseUnitVector);

    SlotRing(uint64_t log2N, uint64_t p, uint64_t e, uint64_t g2, poly baseUnitVector);

    SlotRing(const SlotRing &) = default;
    SlotRing(SlotRing &&) = default;
    ~SlotRing() = default;

    size_t SlotGroupDims() const noexcept;
    size_t SlotGroupLen() const noexcept;
    size_t SlotRank() const noexcept;

    Rotation BlockRotate(size_t slots, size_t blockSize) const;

    /*
     * 计算一个同构，使得每个槽位都按照给定的槽位数量进行移动。
     * 请注意，这里需要进行相对较重的预计算，因此返回的是一个可以重用的函数
     */
    inline Rotation Rotate(size_t slots) const
    {
        return BlockRotate(slots, SlotGroupLen());
    }

    inline Frobenius Frobenius_(size_t iters) const
    {
        return Frobenius(iters, *this);
    }

    inline RawAuto RawAuto_(size_t g1Power, size_t g2Power) const
    {
        return RawAuto(*this, std::make_tuple(g1Power, g2Power));
    }
    uint64_t G1Ord() const noexcept;
    uint64_t G2Ord() const noexcept;
    bool G1GeneratedSlotgroup() const noexcept;

    SubringView R() const noexcept;
    SubringView slot() const noexcept;

    poly ExtractSlotValue(poly x, size_t slot) const;
    poly FromSlotValue(const poly &x, size_t slot) const;
    const poly &SlotOne(size_t slot) const noexcept;
    uint64_t Prime() const noexcept;
    size_t Exponent() const noexcept;
    uint64_t N() const noexcept;
    const uint64_t &IndexMod() const noexcept;

    static bool IsPolySparse(const poly &xPowerN, size_t dPoly);

    bool IsDHalfSparse() const;
    bool IsDSparse() const;

    // 返回一个由x^k生成的subring，其中k = N/2^indexlog2
    SlotRing PowerXSubring(size_t indexLog2) const;
    void InPowerXSubring(const SlotRing &subring, const poly &el, poly &result) const;
    void FromPowerXSubring(const SlotRing &subring, const poly &el, poly &result) const;
    SlotRing ChangeExponent(uint64_t newExp) const;

    friend void TestGAutomorphisms();
};

std::shared_ptr<SlotRing> SmallTestParameters();
std::shared_ptr<SlotRing> SmallP257TestParameters();
std::shared_ptr<SlotRing> P127TestParameters();
std::shared_ptr<SlotRing> P257TestParameters();
void TestGAutomorphisms();
void TestBlockRotate();
void TestRotateNoncyclic();

inline size_t SlotRing::SlotGroupDims() const noexcept
{
    return 1;
}

inline size_t SlotRing::SlotGroupLen() const noexcept
{
    return n;
}

inline size_t SlotRing::SlotRank() const noexcept
{
    return d;
}

inline SlotRing::SubringView SlotRing::R() const noexcept
{
    return { modPe, mod2NCyclotomic };
}

inline SlotRing::SubringView SlotRing::slot() const noexcept
{
    return { modPe, slotModulus };
}

inline uint64_t SlotRing::G1Ord() const noexcept
{
    return N() / 2;
}

inline uint64_t SlotRing::G2Ord() const noexcept
{
    return 2;
}

inline bool SlotRing::G1GeneratedSlotgroup() const noexcept
{
    return slotGroupCyclic;
}

inline uint64_t SlotRing::Prime() const noexcept
{
    return p;
}

inline size_t SlotRing::Exponent() const noexcept
{
    return e;
}

inline const uint64_t &SlotRing::IndexMod() const noexcept
{
    return mod2N;
}
#endif